from sentence_transformers import SentenceTransformer
import pickle
import torch


with open("query5.pkl", "rb") as f:
    query5_list = pickle.load(f)
query5_list=query5_list[:10]
sentences = ["This is an example sentence", "Each sentence is converted"]

model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

# 查询4

import sys
# import pymysql
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QLineEdit, QPushButton, QTextEdit
import json

class QueryWindow(QWidget):
    def __init__(self):
        super().__init__()

        self.init_ui()

    def init_ui(self):
        # 设置窗口标题和大小
        self.setWindowTitle('查询5')
        self.setGeometry(100, 100, 500, 300)

        layout = QVBoxLayout()

        # 创建提示标签
        id_label = QLabel('请输入查询语句:')
        layout.addWidget(id_label)

        # 创建输入框
        self.id_edit = QLineEdit()
        layout.addWidget(self.id_edit)

        # 创建查询按钮
        query_button = QPushButton('查询')
        query_button.clicked.connect(self.query_data)
        layout.addWidget(query_button)

        # 创建用于显示结果的文本框
        self.result_text = QTextEdit()
        self.result_text.setReadOnly(True)
        layout.addWidget(self.result_text)

        self.setLayout(layout)

    def query_data(self):
        input_id = self.id_edit.text()
        if not input_id:
            self.result_text.setText("请输入查询语句")
            return
        try:
            q='what is the population of Kampen?'
            q=input_id
            res=[]
            with torch.no_grad():
                q_embed = model.encode(q)
            for q in query5_list:
                #计算q_embed和q['embed']的余弦相似度
                cos_sim = torch.nn.functional.cosine_similarity(torch.tensor(q_embed), torch.tensor(q['embed']), dim=0)
                item=q.copy()
                #item删除embed字段
                del item['embed']
                item['score']=cos_sim.item()
                res.append(item)
            res.sort(key=lambda x:x['score'],reverse=True)
            self.result_text.setText(f"前5条记录\n{res[0]}\n{res[1]}\n{res[2]}\n{res[3]}\n{res[4]}")
            
        except Exception as err:
            self.result_text.setText(f"查询出错: {err}")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = QueryWindow()
    window.show()
    sys.exit(app.exec_())